import time
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier, plot_tree

warnings.filterwarnings('ignore')

print("正在绘制决策树流程图...")
start = time.time()
test = pd.read_csv('./data/datasets/test.csv')
target = test["predict"]
target = np.array(target)
test = test.drop(["stroke", "id", "predict"], axis=1)
feature_names = test.columns[:10]
data = np.array(test)
# print(data)  # 显示 data 测试数据
# print(target)  # 显示 target 预测结果数据
# print(feature_names)
plt.figure()
plt.axes([0, 0, 1, 1])
clf = DecisionTreeClassifier().fit(data, target)
plot_tree(clf, max_depth=None, feature_names=feature_names, filled=True,
          impurity=True, rotate=True, rounded=True,
          precision=3)
plt.savefig('./data/img/tree_visualization.jpg')
plt.show()
end = time.time()
print("绘图完毕,用时 {} s".format(round(end - start, 3)))
